// --------------------------------------------------------------------------------------------------------------------
// <copyright file="ArmVisualizationAgent.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------

namespace Microsoft.Robotics.Manipulation.Visualization
{
    using System;
    using System.Collections.Generic;
    using System.IO;
    using System.Runtime.Serialization;
    using System.Threading.Tasks;
    using System.Windows.Media;
    using System.Windows.Media.Media3D;

    using Microsoft.Robotics.Geometry;
    using Microsoft.Robotics.Kinematics;
    using Microsoft.Robotics.Manipulation.Runtime;
    using Microsoft.Robotics.Numerics;
    using Microsoft.Robotics.Runtime;
    using Microsoft.Robotics.Visualization;

    using ShoNS.Array;

    /// <summary>
    /// A class that translates manipulation status messages into a concrete scene
    /// for display
    /// </summary>
    [DataContract]
    public class ArmVisualizationAgent : ConsumerAgent<ManipulatorStatusMessage>,
        IControlReceiver<SimulateEETrajectoryMessage>,
        IControlReceiver<SimulateJointTrajectoryMessage>
    {
        /// <summary>
        /// name of the visualization agent
        /// </summary>
        [DataMember]
        private string visualizationAgentName;
        
        /// <summary>
        /// kinematic chain for forward kinematics
        /// </summary>
        [DataMember]
        private KinematicChain armKinematicChain;

        /// <summary>
        /// Maximum number of waypoints to plot out when we visualize a trajectory of
        /// an end effector. One coordinate frame will be plotted out at each waypoint for visualization.
        /// If a given trajectory contains more waypoints than this number,
        /// only this many waypoints will be selected to visualize the given trajectory.
        /// They will be evenly distributed along the trajectory in the visualization process.
        /// </summary>
        [DataMember]
        private int maximumNumberOfSamplingWaypoints;

        /// <summary>
        /// visualization agent
        /// </summary>
        private IAgentHost visualizationAgent;

        /// <summary>
        /// 3d model group
        /// </summary>
        private Model3DGroup model3DGroup;

        /// <summary>
        /// object names
        /// </summary>
        private List<string> linkName;

        /// <summary>
        /// object geometry
        /// </summary>
        private List<Model3D> linkGeometry;
        
        /// <summary>
        /// Initializes a new instance of the <see cref="ArmVisualizationAgent" /> class.
        /// </summary>
        /// <param name="name">name of this agent</param>
        /// <param name="producer">the producer</param>
        /// <param name="visualizationAgent">the visualization agent</param>
        /// <param name="kc">arm kinematic chain</param>
        /// <param name="maximumWaypoints">maximum number of waypoints to visualize for one trajectory</param>
        public ArmVisualizationAgent(string name, IProducer<ManipulatorStatusMessage> producer, IControllable<UpdateRobotMessage> visualizationAgent, KinematicChain kc, int maximumWaypoints = 20)
            : base(name, producer.Name)
        {
            this.visualizationAgentName = visualizationAgent.Name;
            this.armKinematicChain = kc;
            this.maximumNumberOfSamplingWaypoints = maximumWaypoints;
        }

        /// <summary>
        /// Initialization routine
        /// </summary>
        /// <param name="locator">agent locator</param>
        public override void Initialize(AgentLocator locator)
        {
            base.Initialize(locator);
            this.visualizationAgent = this.Locator.GetAgent(this.visualizationAgentName);

            this.linkName = new List<string>();
            this.linkGeometry = new List<Model3D>();
        }

        /// <summary>
        /// When one manipulator status message is received, a new scene with the arm
        /// at the current pose is synthesized.  The new scene is also posted to the
        /// visualizer for display.
        /// </summary>
        /// <param name="message">message received</param>
        public override void Receive(ManipulatorStatusMessage message)
        {
            Matrix4[] transfs = this.armKinematicChain.GetJointTransforms(DoubleArray.From(message.State.JointPositions).Transpose());

            Matrix4[] robotLinkTransforms = new Matrix4[1 + transfs.Length];
            robotLinkTransforms[0] = Matrix4.Identity();
            transfs.CopyTo(robotLinkTransforms, 1);

            this.CreateRobotLinkGeometryGroup(this.armKinematicChain.LinkNames, robotLinkTransforms, this.armKinematicChain.LinkGeometryPaths);
            this.CreateRobotEndEffectorGeometryGroup(this.armKinematicChain.EndEffectorName, robotLinkTransforms[robotLinkTransforms.Length - 1], this.armKinematicChain.EndEffectorGeometryPath);

            if (null != message.State.TcpForceTorque)
            {
                this.CreateRobotForceTorqueRepresentationGeometryGroup(robotLinkTransforms[robotLinkTransforms.Length - 1], message.State.TcpForceTorque, -0.01, 0.001);
            }

            this.CreateRobotCollisionModel(this.armKinematicChain, message.State.JointPositions);

            UpdateRobotMessage updateMsg = new UpdateRobotMessage(this.model3DGroup.Clone()); 
            updateMsg.LinkGroup.Freeze();

            this.visualizationAgent.GetControlPort<UpdateRobotMessage>().Post(updateMsg);
        }

        /// <summary>
        /// When a joint trajectory simulation message is received,
        /// the motion is being synthesized and simulated.
        /// </summary>
        /// <param name="message">message received</param>
        public void ReceiveControl(SimulateEETrajectoryMessage message)
        {
            List<double[]> traj = message.JointTrajectory;

            Model3DGroup frames = new Model3DGroup();

            for (int i = 0; i < traj.Count; i += Math.Max(1, traj.Count / this.maximumNumberOfSamplingWaypoints))
            {
                double[] waypoint = traj[i];

                Matrix4[] transfs = this.armKinematicChain.GetJointTransforms(DoubleArray.From(waypoint).Transpose());

                frames.Children.Add(GeometryManipulator.CreateAFrame(transfs[transfs.Length - 1]));
            }

            UpdateFrameTrajectoryMessage msg = new UpdateFrameTrajectoryMessage(frames.Clone());
            msg.FrameTrajectory.Freeze();
            this.visualizationAgent.GetControlPort<UpdateFrameTrajectoryMessage>().Post(msg);
        }

        /// <summary>
        /// When a joint trajectory simulation message is received,
        /// the motion is being synthesized and simulated.
        /// </summary>
        /// <param name="message">message received</param>
        public void ReceiveControl(SimulateJointTrajectoryMessage message)
        {
            List<double[]> traj = message.JointTrajectory;
            
            // Determine the number of transforms we have
            int numTransforms = 0;
            if (traj.Count > 0)
            {
                Matrix4[] tempTransfs = this.armKinematicChain.GetJointTransforms(DoubleArray.From(traj[0]).Transpose());
                numTransforms = tempTransfs.Length;
            }

            // Allocate memory for robotLinkTransforms apriori
            Matrix4[] robotLinkTransforms = new Matrix4[1 + numTransforms];
            robotLinkTransforms[0] = Matrix4.Identity();

            for (int i = 0; i < traj.Count; ++i)
            {
                double[] waypoint = traj[i];

                Matrix4[] transfs = this.armKinematicChain.GetJointTransforms(DoubleArray.From(waypoint).Transpose());
                transfs.CopyTo(robotLinkTransforms, 1);
                
                this.CreateRobotLinkGeometryGroup(this.armKinematicChain.LinkNames, robotLinkTransforms, this.armKinematicChain.LinkGeometryPaths);
                this.CreateRobotEndEffectorGeometryGroup(this.armKinematicChain.EndEffectorName, robotLinkTransforms[robotLinkTransforms.Length - 1], this.armKinematicChain.EndEffectorGeometryPath);
                
                UpdateRobotMessage updateMsg = new UpdateRobotMessage(this.model3DGroup.Clone());
                updateMsg.IsSimulated = true;
                updateMsg.LinkGroup.Freeze();

                this.visualizationAgent.GetControlPort<UpdateRobotMessage>().Post(updateMsg);

                // In order to visualize the arm along an entire trajectory, we pause/sleep a few millisecond at each waypoint along the trajectory 
                // This is so that one frame gets rendered on the visualizer.
                System.Threading.Thread.Sleep(50);
            }
        }

        /// <summary>
        /// synthesize the scene for the arm at the current pose
        /// </summary>
        /// <param name="names">Link names</param>
        /// <param name="transforms">Transforms for each link</param>
        /// <param name="linkGeometryPaths">Path for each geometry</param>
        public void CreateRobotLinkGeometryGroup(string[] names, Matrix4[] transforms, string[] linkGeometryPaths)
        {
            this.model3DGroup = new Model3DGroup();

            for (int i = 0; i < names.Length; ++i)
            {
                // populate the transformations
                // need to do this transpose, weird
                Matrix4 transf = Matrix4.Transpose(transforms[i]);
                Matrix3D matrix = new Matrix3D(transf[0, 0], transf[0, 1], transf[0, 2], transf[0, 3], transf[1, 0], transf[1, 1], transf[1, 2], transf[1, 3], transf[2, 0], transf[2, 1], transf[2, 2], transf[2, 3], transf[3, 0], transf[3, 1], transf[3, 2], transf[3, 3]);

                MatrixTransform3D transform = new MatrixTransform3D(matrix);

                Model3D model = this.LoadModel(names[i], linkGeometryPaths[i], Colors.DarkOrange);

                model.Transform = transform;

                this.model3DGroup.Children.Add(model);
            }
        }

        /// <summary>
        /// synthesize the scene for the end effector
        /// </summary>
        /// <param name="name">End effector name</param>
        /// <param name="transf">Transforms for the end effector</param>
        /// <param name="path">Path for the geometry</param>
        public void CreateRobotEndEffectorGeometryGroup(string name, Matrix4 transf, string path)
        {
            if (null == this.model3DGroup || null == name || null == path)
            {
                return;
            }

            Model3D model = this.LoadModel(name, path, Colors.DarkBlue);

            transf = Matrix4.Transpose(transf);
            Matrix3D matrix = new Matrix3D(transf[0, 0], transf[0, 1], transf[0, 2], transf[0, 3], transf[1, 0], transf[1, 1], transf[1, 2], transf[1, 3], transf[2, 0], transf[2, 1], transf[2, 2], transf[2, 3], transf[3, 0], transf[3, 1], transf[3, 2], transf[3, 3]);
            MatrixTransform3D transform = new MatrixTransform3D(matrix);
            model.Transform = transform;

            this.model3DGroup.Children.Add(model);
        }

        /// <summary>
        /// Create force torque representation for external force and torque
        /// </summary>
        /// <param name="endEffectorTransf">Transformation for the end effector</param>
        /// <param name="forceTorque">Force and torque information</param>
        /// <param name="scalarForForceTorque">Scalar that is used on raw force torque value for visualization
        /// Value can be modified such that the size of the visualized arrow is in the right proportion compared to
        /// other visualized components, such as the robot.  The default value is negative because the raw data
        /// reported in the Baker system is describing the force and torque the robot is exerting.
        /// Values are empirically chosen.</param>
        /// <param name="scalarForArrowDiameter">Scalar that is used on the total force value for visualization.
        /// Value can be modified such that the size of the visualized arrow is in the right proportion compared to
        /// other visualized components, such as the robot.</param>
        public void CreateRobotForceTorqueRepresentationGeometryGroup(Matrix4 endEffectorTransf, double[] forceTorque, double scalarForForceTorque = -0.01, double scalarForArrowDiameter = 0.001)
        {
            Vector3 origin = new Vector3(endEffectorTransf[0, 3], endEffectorTransf[1, 3], endEffectorTransf[2, 3]);
            Vector3 endEffectorXAxisInWorld = new Vector3(endEffectorTransf[0, 0], endEffectorTransf[1, 0], endEffectorTransf[2, 0]);
            Vector3 endEffectorYAxisInWorld = new Vector3(endEffectorTransf[0, 1], endEffectorTransf[1, 1], endEffectorTransf[2, 1]);
            Vector3 endEffectorZAxisInWorld = new Vector3(endEffectorTransf[0, 2], endEffectorTransf[1, 2], endEffectorTransf[2, 2]);
            Vector3 pointingTowards = origin + scalarForForceTorque * forceTorque[0] * endEffectorXAxisInWorld + scalarForForceTorque * forceTorque[1] * endEffectorYAxisInWorld + scalarForForceTorque * forceTorque[2] * endEffectorZAxisInWorld;

            double totalForce = Math.Sqrt(forceTorque[0] * forceTorque[0] + forceTorque[1] * forceTorque[1] + forceTorque[2] * forceTorque[2]);
            Model3D model = GeometryManipulator.CreateAnArrow(
                new Point3D(origin.X, origin.Y, origin.Z),
                new Point3D(pointingTowards.X, pointingTowards.Y, pointingTowards.Z),
                scalarForArrowDiameter * totalForce,
                Colors.Red);

            this.model3DGroup.Children.Add(model);
        }

        /// <summary>
        /// Create the collision model of the arm
        /// </summary>
        /// <param name="kc">Kinematic chain</param>
        /// <param name="jointValues">Joint values</param>
        public void CreateRobotCollisionModel(KinematicChain kc, double[] jointValues)
        {
            List<GeometryElementGroup3D> collisionModels = kc.GetTransformedCollisionModels(DoubleArray.From(jointValues).Transpose());
            for (int i = 0; i < collisionModels.Count; ++i)
            {
                if (null == collisionModels[i])
                {
                    continue;
                }

                List<GeometryElement3D> currentLink = collisionModels[i].ElementGroup;
                for (int j = 0; j < currentLink.Count; ++j)
                {
                    // In our current implementation, the collision model of each link is a union of spheres.
                    // TODO: revisit this once we have more complex collision model for links
                    SphereElement currentSphere = (SphereElement)currentLink[j];

                    Matrix4 pose = currentSphere.ReferenceFrame * HomogenousTransform.CreateFromTranslation(currentSphere.Location);
                    Model3D model = GeometryManipulator.CreateASphere(new Point3D(pose[0, 3], pose[1, 3], pose[2, 3]), currentSphere.Radius, Colors.Red);

                    this.model3DGroup.Children.Add(model);
                }
            }
        }

        /// <summary>
        /// Load an geometry
        /// </summary>
        /// <param name="name">Name for the geometry</param>
        /// <param name="path">Geometry path</param>
        /// <param name="color">The color of this geometry</param>
        /// <returns>A geometry model</returns>
        public Model3D LoadModel(string name, string path, Color color)
        {
            // check if this model has been loaded or not
            for (int i = 0; i < this.linkName.Count; ++i)
            {
                if (this.linkName[i].Equals(name))
                {
                    return this.linkGeometry[i];
                }
            }

            Model3D geometryModel = GeometryManipulator.Load3DModel(name, path, color);

            // Store the information
            this.linkName.Add(name);
            this.linkGeometry.Add(geometryModel);

            return geometryModel;
        }
    }
}
